import os
import sys
from pathlib import Path
import torch
from paddleocr import PaddleOCR
from PIL import ImageFont, ImageDraw, Image
import numpy as np
import cv2

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # 根目录
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))

from models.common import DetectMultiBackend
from utils.dataloaders import LoadImages
from utils.general import check_img_size, non_max_suppression, scale_boxes
from utils.plots import Annotator
from utils.torch_utils import select_device, smart_inference_mode

# 配置中文字体
CHINESE_FONT_PATH = "C:/Windows/Fonts/simsun.ttc"
FONT_SIZE = 40
FONT_COLOR = (255, 0, 0)  # 红色字体


@smart_inference_mode()
def run_inference(
        source_img_path,
        weights=ROOT / 'best.pt',
        output_dir=ROOT / 'output',
        conf_thres=0.25,
        iou_thres=0.45,
        device=''
):
    """
    车牌识别核心函数
    返回: (处理后的图片路径, 识别结果列表)
    识别结果格式: [{
        'text': 车牌号码,
        'type': '蓝牌'/'绿牌',
        'confidence': 置信度(0-1),
        'position': [x1, y1, x2, y2]
    }, ...]
    """
    # 初始化OCR
    ocr = PaddleOCR(use_angle_cls=True, lang='ch')

    # 加载字体
    try:
        font = ImageFont.truetype(CHINESE_FONT_PATH, FONT_SIZE)
    except:
        font = ImageFont.load_default()

    # 准备输出目录
    output_dir.mkdir(exist_ok=True)
    output_path = output_dir / f'processed_{Path(source_img_path).name}'

    # 加载模型
    device = select_device(device)
    model = DetectMultiBackend(weights, device=device)
    stride, names, pt = model.stride, model.names, model.pt
    imgsz = check_img_size((640, 640), s=stride)

    # 加载图片
    dataset = LoadImages(source_img_path, img_size=imgsz, stride=stride, auto=pt)

    # 存储识别结果
    plate_results = []

    # 处理图片
    for path, im, im0s, _, _ in dataset:
        im = torch.from_numpy(im).to(device)
        im = im.half() if model.fp16 else im.float()
        im /= 255
        if im.ndimension() == 3:
            im = im[None]

        # 推理
        pred = model(im)
        pred = non_max_suppression(pred[0][1], conf_thres, iou_thres, max_det=1000)

        # 处理结果
        im0 = im0s.copy()
        annotator = Annotator(im0, line_width=3, example=str(names))
        texts_to_draw = []

        if len(pred[0]):
            det = pred[0]
            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()

            for *xyxy, conf, cls in reversed(det):
                if int(cls) in (0, 1):  # 只处理车牌(0:蓝牌, 1:绿牌)
                    x1, y1, x2, y2 = map(int, xyxy)
                    plate_img = im0[y1:y2, x1:x2]

                    # OCR识别
                    ocr_res = ocr.ocr(plate_img, cls=True)
                    plate_text = ''.join([line[1][0] for line in ocr_res[0]]) if ocr_res else ''

                    # 获取车牌类型
                    plate_type = "蓝牌" if int(cls) == 0 else "绿牌"

                    # 保存结果
                    plate_results.append({
                        'text': plate_text,
                        'type': plate_type,
                        'confidence': float(conf),
                        'position': [x1, y1, x2, y2]
                    })

                    # 准备绘制文本 (类型:号码)
                    display_text = f"{plate_text}"
                    texts_to_draw.append((display_text, (x1, y2 + 5)))

                # 绘制检测框
                label = f'{names[int(cls)]} {conf:.2f}'
                annotator.box_label(xyxy, label)

        # 绘制OCR识别文本
        im0 = annotator.result()
        if texts_to_draw:
            img_pil = Image.fromarray(cv2.cvtColor(im0, cv2.COLOR_BGR2RGB))
            draw = ImageDraw.Draw(img_pil)
            for txt, (tx, ty) in texts_to_draw:
                draw.text((tx, ty), txt, font=font, fill=FONT_COLOR)
            im0 = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)

        # 保存结果图片
        cv2.imwrite(str(output_path), im0)

    return str(output_path), plate_results
